/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "test_dtls.h"

#define BOOST_TEST_DYN_LINK
#ifdef STAND_ALONE_TEST
#define BOOST_TEST_MODULE dtlstest
#endif
#include <boost/test/unit_test.hpp>

using namespace aiengine;

BOOST_FIXTURE_TEST_SUITE(dtls_test_suite, StackDTLStest)

BOOST_AUTO_TEST_CASE (test01)
{
	Packet packet;

	BOOST_CHECK(dtls->getTotalBytes() == 0);
	BOOST_CHECK(dtls->getTotalPackets() == 0);
	BOOST_CHECK(dtls->getTotalValidPackets() == 0);
	BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);
	BOOST_CHECK(dtls->processPacket(packet) == true);

	CounterMap c = dtls->getCounters();

	BOOST_CHECK(dtls->getId() == 0x00);
	BOOST_CHECK(dtls->getHeaderSize() == sizeof(dtls_header));
}

BOOST_AUTO_TEST_CASE (test02) // Client Hello
{
	Packet packet("../dtls/packets/packet01.pcap");

	inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        auto info = flow->getDTLSInfo();

        BOOST_CHECK(info != nullptr);

	BOOST_CHECK(dtls->getTotalBytes() == 119);
	BOOST_CHECK(dtls->getTotalPackets() == 1);
	BOOST_CHECK(dtls->getTotalValidPackets() == 1);
	BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

	// DTLS info values
	BOOST_CHECK(info->isEncrypted() == false);
	BOOST_CHECK(info->getTotalDataPdus() == 0);
	BOOST_CHECK(info->getVersion() == 0x0100);
	BOOST_CHECK(info->getEpoch() == 0);

	// DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 0);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 1);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 1);

	CounterMap c = dtls->getCounters();
}

BOOST_AUTO_TEST_CASE (test03) // Hello verify request
{
        Packet packet("../dtls/packets/packet02.pcap");

        inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        BOOST_CHECK(dtls->getTotalBytes() == 28);
        BOOST_CHECK(dtls->getTotalPackets() == 1);
        BOOST_CHECK(dtls->getTotalValidPackets() == 1);
        BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 0);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 0);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 1);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 1);
}

BOOST_AUTO_TEST_CASE (test04) // server hello, cert and server hello done
{
        Packet packet("../dtls/packets/packet03.pcap");

        inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        BOOST_CHECK(dtls->getTotalBytes() == 977);
        BOOST_CHECK(dtls->getTotalPackets() == 1);
        BOOST_CHECK(dtls->getTotalValidPackets() == 1);
        BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 3);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 0);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 0);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 1);
        BOOST_CHECK(dtls->getTotalCertificates() == 1);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 1);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 3);
}

BOOST_AUTO_TEST_CASE (test05) // client key exchange, change cipher, encrypted handshake
{
        Packet packet("../dtls/packets/packet04.pcap");

        inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        auto info = flow->getDTLSInfo();

        BOOST_CHECK(info != nullptr);

	// DTLS info values
	BOOST_CHECK(info->isEncrypted() == true);
	BOOST_CHECK(info->getTotalDataPdus() == 0);
	BOOST_CHECK(info->getVersion() == 0x0100);
        BOOST_CHECK(info->getEpoch() == 0);

        BOOST_CHECK(dtls->getTotalBytes() == 246);
        BOOST_CHECK(dtls->getTotalPackets() == 1);
        BOOST_CHECK(dtls->getTotalValidPackets() == 1);
        BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 1);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 0);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 1);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 3);

	Json j;

	dtls->statistics(j, 5);

	BOOST_CHECK(j["alerts"] == 0);
	BOOST_CHECK(j["bytes"] == 246);
	BOOST_CHECK(j["change cipher specs"] == 1);
	BOOST_CHECK(j["encrypt handshakes"] == 1);
	BOOST_CHECK(j["handshakes"] == 1);
	BOOST_CHECK(j["types"]["client hellos"] == 0);
	BOOST_CHECK(j["types"]["client keys"] == 1);
	BOOST_CHECK(j["types"]["records"] == 3);
}

BOOST_AUTO_TEST_CASE (test06) // new session ticket, change cipher, encrypted handshake
{
        Packet packet("../dtls/packets/packet05.pcap");

        inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        auto info = flow->getDTLSInfo();

        BOOST_CHECK(info != nullptr);

        // DTLS info values
        BOOST_CHECK(info->isEncrypted() == true);
        BOOST_CHECK(info->getTotalDataPdus() == 0);
        BOOST_CHECK(info->getVersion() == 0xfeff);

        BOOST_CHECK(dtls->getTotalBytes() == 282);
        BOOST_CHECK(dtls->getTotalPackets() == 1);
        BOOST_CHECK(dtls->getTotalValidPackets() == 1);
        BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 1);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 0);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 1);
        BOOST_CHECK(dtls->getTotalRecords() == 3);
}

BOOST_AUTO_TEST_CASE (test07) // new session ticket, change cipher, encrypted handshake
{
        Packet packet("../dtls/packets/packet07.pcap");

        inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

        auto info = flow->getDTLSInfo();

        BOOST_CHECK(info != nullptr);

        // DTLS info values
        BOOST_CHECK(info->isEncrypted() == false);
        BOOST_CHECK(info->getTotalDataPdus() == 1);
        BOOST_CHECK(info->getVersion() == 0xfefd);
        BOOST_CHECK(info->getEpoch() == 1);

        BOOST_CHECK(dtls->getTotalBytes() == 93);
        BOOST_CHECK(dtls->getTotalPackets() == 1);
        BOOST_CHECK(dtls->getTotalValidPackets() == 1);
        BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 0);
        BOOST_CHECK(dtls->getTotalDatas() == 1);

        BOOST_CHECK(dtls->getTotalClientHellos() == 0);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 1);
}

BOOST_AUTO_TEST_SUITE_END()

BOOST_FIXTURE_TEST_SUITE(ipv6_dtls_test_suite, StackIPv6DTLStest)

BOOST_AUTO_TEST_CASE (test01)
{
	Packet packet("../dtls/packets/packet06.pcap");

	inject(packet);

        Flow *flow = dtls->getCurrentFlow();

        BOOST_CHECK(flow != nullptr);

	BOOST_CHECK(dtls->getTotalBytes() == 421);
	BOOST_CHECK(dtls->getTotalPackets() == 1);
	BOOST_CHECK(dtls->getTotalValidPackets() == 1);
	BOOST_CHECK(dtls->getTotalInvalidPackets() == 0);

	Json j;

	dtls->statistics(j, 5);

        // DTLS counters check
        BOOST_CHECK(dtls->getTotalHandshakes() == 1);
        BOOST_CHECK(dtls->getTotalEncryptedHandshakes() == 0);
        BOOST_CHECK(dtls->getTotalAlerts() == 0);
        BOOST_CHECK(dtls->getTotalChangeCipherSpecs() == 0);
        BOOST_CHECK(dtls->getTotalDatas() == 0);

        BOOST_CHECK(dtls->getTotalClientHellos() == 1);
        BOOST_CHECK(dtls->getTotalHelloVerifyRequests() == 0);
        BOOST_CHECK(dtls->getTotalServerHellos() == 0);
        BOOST_CHECK(dtls->getTotalCertificates() == 0);
        BOOST_CHECK(dtls->getTotalCertificateRequests() == 0);
        BOOST_CHECK(dtls->getTotalCertificateVerifies() == 0);
        BOOST_CHECK(dtls->getTotalServerDones() == 0);
        BOOST_CHECK(dtls->getTotalServerKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalClientKeyExchanges() == 0);
        BOOST_CHECK(dtls->getTotalHandshakeFinishes() == 0);
        BOOST_CHECK(dtls->getTotalNewSessionTickets() == 0);
        BOOST_CHECK(dtls->getTotalRecords() == 1);

	BOOST_CHECK(j["alerts"] == 0);
	BOOST_CHECK(j["bytes"] == 421);
	BOOST_CHECK(j["change cipher specs"] == 0);
	BOOST_CHECK(j["encrypt handshakes"] == 0);
	BOOST_CHECK(j["handshakes"] == 1);
	BOOST_CHECK(j["types"]["client hellos"] == 1);
	BOOST_CHECK(j["types"]["client keys"] == 0);
	BOOST_CHECK(j["types"]["records"] == 1);

	dtls->decreaseAllocatedMemory(10);
}

BOOST_AUTO_TEST_SUITE_END()
